#include <stdio.h>
#include <string.h>
#include "Csm.h"
#include "CryIf.h"
#include "Crypto.h"
#include "Os.h"

// --------------------------------------------------------------------------------
// * Test Data Storage
// --------------------------------------------------------------------------------
#define MESSAGE_BUFFER_SIZE 2048
static uint8_t private_key[32];
static uint8_t public_key[64];
static uint8_t plain_message[MESSAGE_BUFFER_SIZE];
static uint8_t encrypted_message[MESSAGE_BUFFER_SIZE];
static uint8_t signature[64];
// static uint8_t ecc_k[32];


// --------------------------------------------------------------------------------
// * Helper Functions
// --------------------------------------------------------------------------------
int hex_char_to_int(char c) {
	if (c >= 'a' && c <= 'z') {
		return c - 'a' + 10;
	}
	if (c >= 'A' && c <= 'Z') {
		return c - 'A' + 10;
	}
	if (c >= '0' && c <= '9') {
		return c - '0';
	}
	return 0;
}
void hex_str_to_byte_array(uint8_t* buf, const char* hex_str) {
	uint32_t slen = strlen(hex_str);
	for (uint32_t i = 0; i < slen; i += 2) {
		int v1 = hex_char_to_int(hex_str[i]);
		int v0 = hex_char_to_int(hex_str[i + 1]);
		buf[i / 2] = v1 * 16 + v0;
	}
}


// --------------------------------------------------------------------------------
// Tests
// --------------------------------------------------------------------------------
void run_sm2_test_setup() {
	// Init Crypto Stack
	Crypto_Init(NULL_PTR);
	CryIf_Init(NULL_PTR);
	Csm_Init(NULL_PTR);

}

void run_sm2_signature_generate_verify_test() {

	const char* user_id_cstr = "1234567812345678";
	const char* private_key_hex = "59A408B5ED69E2EBE8DC3DD4621453DA8B0DEBF294190ACD9CDB452E90B7763D";
	const char* public_key_hex = "2EB5019334A21AE618570009491C50B3C17D6E6B6B9CDEAF7E924B5004852252CE93DF59CFB607B4ED46A654341F4636C12333C3FA13E28F5B7404A30C0AFA84";
	const char* input_data_hex = "7C93AA2CB289D3639453CE12408F81295333528D0481163C76C92C2ED5A28F4846E63BC46A17AD1F894D48F71ADD661C00529B961C6D86E7B4B0E60246F0D814100762ACA15F89C78126C54AC7CBF742391B65C70C9D827EB99C33AE14F6D2660A80D3CDB372D142D4C9FD158D7722F81C7441213AF3159C53417D10692FD6043A792CA66CBD014AE0CEC76918F774F4EBA750D68F88D5413D54AE08C9B1A8F7D7A61B9C59E163A14D9BF99F472013C2E6E2E76FD885B682B135D90F2F4945E99489AA807FDD8D24745B2623C253BF55C3DE6A91A612A5979090EF7D804D0451EA5A2B1F23A3E97204246FB3FEFECBF1BCEFA42F447DA263289C84856A3ED5DCBEBBEA1C9FF986C0B44508EB8B04309F7228A8BE63D3212DAA17C7E8E2EE7644C5F31B7ACD7AF4E98C03BAB754B7CC742641BBA73E397D3B828B1CA20C749727BD3BDC45ACEEC18B82D476AA27ED36A61295AD914E0A69ED24726FDF2DF63A1D7EA8997558CA9DD9C44E15F1F40B5061A29FA7E43BD87F47069D420675C84B8CA5037C20BA1759CC3F8224392B3BF586BB429A38830CB8191550312ADBBF12B18F537D3FC5F40EF6F4D72A2F4245AE3F30F2654EB4C5C68F527612D976D901B7BD8AEB5E8E9071B91C22A3039531A3BCE0FF1D8DF55542AF9D2D147552E5771C3879E4407B0E86FD3012BDFA725E3BDD4233F8EA2A914D0006DC93D832DE11C7060EB2D9F90CAE502672EF8EAC2A31BD3B3594230CAA92DB6D0066FFA6C2A8E20EF56E2695EE5BC017D18A4F00A68F4C260BD672E2C5B8DF4CBB70C64A5E03032E50AB9E997A0F1FB0C498328EA6E7856A3FFBBA5F467745DBE1460EA785C61B4B07FE39E661A5B8F62C145E534EFF7D06CC343A6D49DF0CB2057D41897C024A1FBCA4857553ED498DEACF1923F0A2A8155FAB0AC5CB2C7837170E0871DF4895C1D0032DF7881F0C5F2BD493223F73BA21465FE34344FF705C88F5A7CE9D997E21CD365D21BB59F5B6B8A89C49EE7A8C68A93F43CDC69CF7A9BC61EB3225648E5B4D90AB136D8259D3B0606FC0B48F6ACC0FE232323F21D137E8C906BE5FDCC73401C2234CC4955ED9EA3D1906605DAC788EFF8A8649E8CE0B887E58164DB6FF5C23A95177E62D504D6C158A3F39F40587A317FDD63E1DC698222CA1CF5901B9F8BFE2E5B373DD2EBCCCA7FD8C1D58B9DF5B7DBEC250098144DCC759563FCA64674D75A6524F6F7E8DAD5870953C6DD11E0FD73680C837908FF07A95D891C151D8F4F5C84C26ECD672D2E5B390E335CBF25855C23B15BC3448BD916EF21528A457C3795F985AF1867685096504C87B60E7574E04A6E0E546588F471E3B77B2D4690030298D67FDEA7C52CE2FEC54AEFB4FFB5981AB44CCD196C2E0192053FC6A01EB372E61FBA5EAE26A4F2E32D295DADF189FEDD14A1F03C568B3D12C09E927";

	printf("\n================ SM2 Sign Generate Test ================\n\n");
	run_sm2_test_setup();

	// Set Private Key
	hex_str_to_byte_array(private_key, private_key_hex);
	Csm_KeyElementSet(CsmKey_SM2_PrivateKey, 0, private_key, 32);
	Csm_KeySetValid(CsmKey_SM2_PrivateKey);

	// Set Public Key
	hex_str_to_byte_array(public_key, public_key_hex);
	Csm_KeyElementSet(CsmKey_SM2_PublicKey, 0, public_key, 64);
	Csm_KeySetValid(CsmKey_SM2_PublicKey);

	// Set User Id
	Csm_KeyElementSet(CsmKey_SM2_UserId, 0, (uint8_t*) user_id_cstr, strlen(user_id_cstr) + 1);
	Csm_KeySetValid(CsmKey_SM2_UserId);

	// Fix a k, so the result can be reproduced
	//bigint_set_from_hex_str(ctx->k, "8DF579B422980688132F7726BB7716C6589DE07E97169589BBBE4B90ED43977B");
	
	// Source data in bytes
	uint32_t plain_message_bytes = strlen(input_data_hex) / 2;
	hex_str_to_byte_array(plain_message, input_data_hex);

	// Signature Generate
	uint32_t signature_bytes = 64;
	Std_ReturnType sign_gen_res = Csm_SignatureGenerate(
		CsmJob_SM2_SignGenerateSync,
		CRYPTO_OPERATIONMODE_SINGLECALL,
		plain_message,
		plain_message_bytes,
		signature,
		&signature_bytes
	);
	if (sign_gen_res != E_OK) {
		printf("[X] Signature Generate Failed: %u\n", sign_gen_res);
	} else {
		printf("[√] Signature Generate Success: (%lu bytes)\n", signature_bytes);
		printf("    r = ");
		for (uint32_t i = 0; i < 32; ++i) {
			printf("%02X", signature[i]);
		}
		printf("\n");
		printf("    s = ");
		for (uint32_t i = 0; i < 32; ++i) {
			printf("%02X", signature[i + 32]);
		}
		printf("\n");
	}

	Schedule();
	printf("\n================ SM2 Sign Verify Test ================\n\n");
	uint8_t verify_result = 0;
	Csm_SignatureVerify(
		CsmJob_SM2_SignVerifySync,
		CRYPTO_OPERATIONMODE_SINGLECALL,
		plain_message,
		plain_message_bytes,
		signature,
		signature_bytes,
		&verify_result
	);
	if (verify_result == E_OK) {
		printf("[√] Signature Verify Pass\n");
	} else {
		printf("[x] Signature Verify Failed, error_code: %d\n", verify_result);
	}
}


void run_sm2_data_encrypt_decrypt_test() {
	const char* private_key_hex = "59A408B5ED69E2EBE8DC3DD4621453DA8B0DEBF294190ACD9CDB452E90B7763D";
	const char* public_key_hex = "2EB5019334A21AE618570009491C50B3C17D6E6B6B9CDEAF7E924B5004852252CE93DF59CFB607B4ED46A654341F4636C12333C3FA13E28F5B7404A30C0AFA84";
	const char* input_data_hex = "966397a223d8521deec7a9c256ba70b53e313fb056e962a35738561eb5b8f4cc322c2433c9ddbcc1b3d13a44a979808db5afc6b43bf43e90fca3ce9f294fce82";

	printf("================ SM2 Encryption Test ================\n\n");
	run_sm2_test_setup();

	// Set Public Key
	hex_str_to_byte_array(public_key, public_key_hex);
	Csm_KeyElementSet(CsmKey_SM2_PublicKey, 0, public_key, 64);
	Csm_KeySetValid(CsmKey_SM2_PublicKey);
	
	// Source data in bytes
	uint32_t plain_message_bytes = strlen(input_data_hex) / 2;
	hex_str_to_byte_array(plain_message, input_data_hex);

	// Encrypt Data
	uint32_t encrypted_message_bytes = 0;
	Std_ReturnType encrypt_res = Csm_Encrypt(
		CsmJob_SM2_DataEncryptSync,
		CRYPTO_OPERATIONMODE_SINGLECALL,
		plain_message,
		plain_message_bytes,
		encrypted_message,
		&encrypted_message_bytes
	);
	if (encrypt_res != E_OK) {
		printf("[X] Encryption Failed: %u\n", encrypt_res);
	} else {
		printf("[√] Encryption Success: (%lu bytes)\n", encrypted_message_bytes);
		for (uint32_t i = 0; i < encrypted_message_bytes; ++i) {
			printf("%02X", encrypted_message[i]);
		}
		printf("\n");
	}

	Schedule();
	printf("\n================ SM2 Decryption Test ================\n\n");
	// Set Private Key
	hex_str_to_byte_array(private_key, private_key_hex);
	Csm_KeyElementSet(CsmKey_SM2_PrivateKey, 0, private_key, 32);
	Csm_KeySetValid(CsmKey_SM2_PrivateKey);

	// Clear result buffer.
	plain_message_bytes = 0;
	memset(plain_message, 0, MESSAGE_BUFFER_SIZE);

	// Decrypt
	Std_ReturnType decrypt_res = Csm_Decrypt(
		CsmJob_SM2_DataDecryptSync,
		CRYPTO_OPERATIONMODE_SINGLECALL,
		encrypted_message,
		encrypted_message_bytes,
		plain_message,
		&plain_message_bytes
	);
	if (decrypt_res != E_OK) {
		printf("[X] Decryption Failed: %u\n", decrypt_res);
	} else {
		printf("[√] Decryption Success: (%lu bytes)\n", plain_message_bytes);
		for (uint32_t i = 0; i < plain_message_bytes; ++i) {
			printf("%02X", plain_message[i]);
		}
		printf("\n");
	}
}


void run_sm2_integration_test() {
	run_sm2_data_encrypt_decrypt_test();
	run_sm2_signature_generate_verify_test();
}


